﻿using System;
using System.IO;
using System.Security.Cryptography;
using MBNCSUtil;

namespace Felbot.Protocol {
    /// <summary>This handles all Warden related "magic", along with its two helper classes, SimpleCrypto and WardonRandom.</summary>
    public class WardenModule {
        /// <summary>Length of the current module.</summary>
        private int m_Length = 0;
        /// <summary>Position (request 0x01) in the current module.</summary>
        private int m_Position = 0;
        /// <summary>MD5 hash of module (aka the .mod file name).</summary>
        private byte[] m_Name;
        /// <summary>The decryption key.</summary>
        private byte[] m_Decr;
        /// <summary>The data gathered in request 0x01.</summary>
        private byte[] m_Data;
        /// <summary>The data decrypted with the decryption key.</summary>
        //private byte[] m_DecrData;

        /// <summary>The WardenRandom class used to generate SimpleCrypto decryption and encryption objects.</summary>
        private WardenRandom m_Random;
        /// <summary>The SimpleCrypto used to decrypt (key in) packets.</summary>
        private SimpleCrypto m_Decryptor;
        /// <summary>The SimpleCrypto used to encrypt (key out) packets.</summary>
        private SimpleCrypto m_Encryptor;

        /// <summary>Object to perform MD5 hashing.</summary>
        private MD5CryptoServiceProvider m_Md5 = new MD5CryptoServiceProvider();

        /// <summary>Called when the Warden packet is recieved, this function does all the work of Warden.</summary>
        /// <param name="a_bEncrypted">The packet, excluding the BNET 4-byte header.</param>
        /// <param name="a_sHashPath">The path to the Hashes folder, in which are the required hashes in sub-folders.</param>
        public byte[] RecieveWarden(byte[] encrypted, string hashPath) {
            DataReader decryptedPacketReader = new DataReader(m_Decryptor.Crypt(encrypted));
            // Console.WriteLine("[WARDEN] Recieved Decrypted:" + Environment.NewLine +
            //     m_drDecryptedPacket.ToString());
            byte request = decryptedPacketReader.ReadByte();
            DataBuffer decryptedPacketBuffer = new DataBuffer();
            // Console.WriteLine("[WARDEN] Request ID 0x" + m_bRequest.ToString("X"));
            switch (request) {
                case 0x00: //request 0x00: WID_CURRENTINFO
                    m_Name = decryptedPacketReader.ReadByteArray(16);
                    m_Decr = decryptedPacketReader.ReadByteArray(16);
                    m_Length = decryptedPacketReader.ReadInt32();
                    m_Position = 0;
                    m_Data = new byte[m_Length];

                    //if (File.Exists(a_sHashPath + "Warden\\" + new string(new BinaryReader(new MemoryStream(g_bName)).ReadChars(16)) + ".mod"))
                    //{
                    decryptedPacketBuffer.InsertByte(0x01);
                    //}
                    //else
                    //{
                    //    m_dbDecryptedPacket.InsertByte(0x00);
                    //}
                    break;

                case 0x01: //request 0x01: WID_CURRENTMODULE
                    short m_shLength = decryptedPacketReader.ReadInt16();
                    decryptedPacketReader.ReadByteArray(m_shLength).CopyTo(m_Data, m_Position);
                    m_Position += m_shLength;

                    if (m_Length > m_Position)
                        return null;

                    #region Initial Validation
                    //if (md5.ComputeHash(g_bData) == g_bName)
                    //{
                    decryptedPacketBuffer.InsertByte(0x01);
                    //}
                    //else
                    //{
                    //    m_dbDecryptedPacket.InsertByte(0x00);
                    //}
                    #endregion

                    #region Decrypt, Revalidate, Decompress, Prepare
                    ///*
                    // * THE FOLLOWING CODE IS UNUSED AS CURRENTLY
                    // * IT IS NOT REQUIRED WHEN RESPONDING TO
                    // * REQUEST 0x02.  THEREFORE, I HAVE REMOVED
                    // * ALL OBJECTS IT USES (QUITE A FEW) AND
                    // * COMMENTED IT OUT, LEAVING COMMENTS AS TO
                    // * WHAT OBJECTS ARE NEEDED FOR IT TO WORK,
                    // * JUST IN CASE FUTURE WARDEN MODULES ARE
                    // * NEEDED TO DECRYPT, REVALIDATE, DECOMPRESS,
                    // * AND PREPARE.
                    // */

                    //#region Decryption
                    //// uses the SimpleCrypto object used in other places.
                    //g_bDecrData = new SimpleCrypto(g_bDecr).Crypt(g_bData);
                    //DataReader m_drDecryptedData = new DataReader(g_bDecrData);
                    //#endregion

                    //#region Second Validation
                    //// uses the BigInteger class copiable directly from MBNCSUtil
                    //// (it's internal, so you have to copy it local for this part to work).
                    //int m_iLengthInnerData = m_drDecryptedData.ReadInt32();
                    //byte[] m_bInnerData = m_drDecryptedData.ReadByteArray(0x108);
                    //if (m_drDecryptedData.ReadDwordString(0x00) != "SIGN")
                    //    throw new Exception("Decryption failed. \"SIGN\" not found.");
                    //byte[] m_bSignature = m_drDecryptedData.ReadByteArray(0x100);
                    //BigInteger m_biPower = new BigInteger(new byte[4] { 0x01, 0x00, 0x01, 0x00});
                    //BigInteger m_biMod =
                    //    new BigInteger(
                    //        new byte[128] {
                    //            0x6B, 0xCE, 0xF5, 0x2D, 0x2A, 0x7D, 0x7A, 0x67,
                    //            0x21, 0x21, 0x84, 0xC9, 0xBC, 0x25, 0xC7, 0xBC,
                    //            0xDF, 0x3D, 0x8F, 0xD9, 0x47, 0xBC, 0x45, 0x48,
                    //            0x8B, 0x22, 0x85, 0x3B, 0xC5, 0xC1, 0xF4, 0xF5,
                    //            0x3C, 0x0C, 0x49, 0xBB, 0x56, 0xE0, 0x3D, 0xBC,
                    //            0xA2, 0xD2, 0x35, 0xC1, 0xF0, 0x74, 0x2E, 0x15,
                    //            0x5A, 0x06, 0x8A, 0x68, 0x01, 0x9E, 0x60, 0x17,
                    //            0x70, 0x8B, 0xBD, 0xF8, 0xD5, 0xF9, 0x3A, 0xD3,
                    //            0x25, 0xB2, 0x66, 0x92, 0xBA, 0x43, 0x8A, 0x81,
                    //            0x52, 0x0F, 0x64, 0x98, 0xFF, 0x60, 0x37, 0xAF,
                    //            0xB4, 0x11, 0x8C, 0xF9, 0x2E, 0xC5, 0xEE, 0xCA,
                    //            0xB4, 0x41, 0x60, 0x3C, 0x7D, 0x02, 0xAF, 0xA1,
                    //            0x2B, 0x9B, 0x22, 0x4B, 0x3B, 0xFC, 0xD2, 0x5D,
                    //            0x73, 0xE9, 0x29, 0x34, 0x91, 0x85, 0x93, 0x4C,
                    //            0xBE, 0xBE, 0x73, 0xA9, 0xD2, 0x3B, 0x27, 0x7A,
                    //            0x47, 0x76, 0xEC, 0xB0, 0x28, 0xC9, 0xC1, 0xDA,
                    //            0xEE, 0xAA, 0xB3, 0x96, 0x9C, 0x1E, 0xF5, 0x6B,
                    //            0xF6, 0x64, 0xD8, 0x94, 0x2E, 0xF1, 0xF7, 0x14,
                    //            0x5F, 0xA0, 0xF1, 0xA3, 0xB9, 0xB1, 0xAA, 0x58,
                    //            0x97, 0xDC, 0x09, 0x17, 0x0C, 0x04, 0xD3, 0x8E,
                    //            0x02, 0x2C, 0x83, 0x8A, 0xD6, 0xAF, 0x7C, 0xFE,
                    //            0x83, 0x33, 0xC6, 0xA8, 0xC3, 0x84, 0xEF, 0x29,
                    //            0x06, 0xA9, 0xB7, 0x2D, 0x06, 0x0B, 0x0D, 0x6F,
                    //            0x70, 0x9E, 0x34, 0xA6, 0xC7, 0x31, 0xBE, 0x56,
                    //            0xDE, 0xDD, 0x02, 0x92, 0xF8, 0xA0, 0x58, 0x0B,
                    //            0xFC, 0xFA, 0xBA, 0x49, 0xB4, 0x48, 0xDB, 0xEC,
                    //            0x25, 0xF3, 0x18, 0x8F, 0x2D, 0xB3, 0xC0, 0xB8,
                    //            0xDD, 0xBC, 0xD6, 0xAA, 0xA6, 0xDB, 0x6F, 0x7D,
                    //            0x7D, 0x25, 0xA6, 0xCD, 0x39, 0x6D, 0xDA, 0x76,
                    //            0x0C, 0x79, 0xBF, 0x48, 0x25, 0xFC, 0x2D, 0xC5,
                    //            0xFA, 0x53, 0x9B, 0x4D, 0x60, 0xF4, 0xEF, 0xC7,
                    //            0xEA, 0xAC, 0xA1, 0x7B, 0x03, 0xF4, 0xAF, 0xC7
                    //        });
                    //byte[] m_bResult = new BigInteger(m_bSignature).ModPow(m_biPower, m_biMod).GetBytes();
                    //byte[] m_bProperResult = new byte[0x100];
                    //for (int i = 0; i < 0x100; i++)
                    //    m_bProperResult[i] = 0xBB;
                    //byte[] m_bShaThis = new byte[0x10C];
                    //new BigInteger(m_iLengthInnerData).GetBytes().CopyTo(m_bShaThis, 0);
                    //m_bInnerData.CopyTo(m_bShaThis, 4);
                    //WardenSHA1.hash(m_bShaThis).CopyTo(m_bProperResult, 0);
                    //for (int i = 0; i < 0x100; i++)
                    //    if (m_bResult[i] != m_bProperResult[i])
                    //        throw new Exception("Decryption failed. RSA check failed.");
                    //#endregion

                    //#region Decompression
                    //// this uses a trimmed ZLIB.NET include, very big...
                    //// using Felbot.Warden.Zlib;
                    //MemoryStream outStream = new MemoryStream(m_bInnerData);
                    //MemoryStream inStream = new MemoryStream();
                    //ZOutputStream outZStream = new ZOutputStream(outStream);
                    //try
                    //{
                    //    byte[] buffer = new byte[2000];
                    //    int len;
                    //    while ((len = inStream.Read(buffer, 0, 2000)) > 0)
                    //    {
                    //        outZStream.Write(buffer, 0, len);
                    //    }
                    //    outZStream.Flush();
                    //}
                    //finally
                    //{
                    //    outZStream.Close();
                    //    outStream.Close();
                    //    inStream.Close();
                    //}
                    //byte[] m_bInnerDataInflated = outStream.ToArray();
                    //#endregion

                    //#region Preparation
                    //// not yet added, I got lazy... here's how to prepare it:
                    //// http://www.skullsecurity.org/wiki/index.php/Warden_Modules#Preperation
                    //#endregion

                    //#region Storage
                    //// this takes the final data and stores it in a mod file...
                    //File.WriteAllBytes(a_sHashPath + "Warden\\" + new string(new BinaryReader(new MemoryStream(g_bName)).ReadChars(16)) + ".mod", m_bInnerDataPrepared);
                    //#endregion
                    #endregion
                    break;

                case 0x02: //request 0x02: WID_VALIDATECHAT
                    MemoryStream fileData = new MemoryStream(File.ReadAllBytes(hashPath + "StarCraft.exe"));
                    byte eventID = decryptedPacketReader.ReadByte();
                    // Console.WriteLine("[WARDEN] Event: 0x" + m_bEventID.ToString("X"));
                    int iterations = (decryptedPacketReader.Length - 3) / 7;
                    int messageLength = 0;
                    int[] command = new int[iterations];
                    int[] address = new int[iterations];
                    int[] length = new int[iterations];
                    byte[][] value = new byte[iterations][];
                    for (int i = 0; i < iterations; i++) {
                        command[i] = decryptedPacketReader.ReadInt16();
                        address[i] = decryptedPacketReader.ReadInt32();
                        length[i] = decryptedPacketReader.ReadByte();
                        // Console.WriteLine("[WARDEN] Command: 0x" + m_iCommand[i].ToString("X") +
                        //                   "  Address: 0x" + m_iAddress[i].ToString("X") +
                        //                   "  Length: 0x" + m_iLength[i].ToString("X"));

                        value[i] = new byte[length[i] + 1];
                        fileData.Seek(address[i] - 0x400000, SeekOrigin.Begin);
                        fileData.Read(value[i], 1, length[i]);
                        messageLength += length[i] + 1;
                    }

                    int checksum = 0;
                    const int address1 = 0x00497FB0;
                    const int address2 = 0x0049C33D;
                    const int address3 = 0x004A2FF7;
                    unchecked {
                        if (address[0] == address1 && address[1] == address2 && address[2] == address3)
                            checksum = (int) 0x193E73E8;
                        else if (address[0] == address2 && address[1] == address1 && address[2] == address3)
                            checksum = (int) 0xD6557DEF;
                        else if (address[0] == address1 && address[1] == address3 && address[2] == address2)
                            checksum = (int) 0x2183172A;
                        else if (address[0] == address2 && address[1] == address3 && address[2] == address1)
                            checksum = (int) 0xCA841860;
                        else if (address[0] == address3 && address[1] == address2 && address[2] == address1)
                            checksum = (int) 0x9F2AD2C3;
                        else if (address[0] == address3 && address[1] == address1 && address[2] == address2)
                            checksum = (int) 0xC04CF757;
                        else
                            throw new NotSupportedException("Unknown Warden addresses.");
                    }

                    decryptedPacketBuffer.InsertByte(0x02); //Request ID
                    decryptedPacketBuffer.InsertInt16(short.Parse(messageLength.ToString())); //Message Length
                    decryptedPacketBuffer.InsertInt32(checksum); //Checksum
                    for (int i = 0; i < iterations; i++) {
                        decryptedPacketBuffer.InsertByteArray(value[i]); //Value[i]
                    }
                    break;

                default:
                    throw new NotSupportedException("Unknown Warden request.");
            }
            // Console.WriteLine("[WARDEN] Decrypted:" + Environment.NewLine +
            //     m_dbDecryptedPacket.ToString());
            byte[] m_bEncrypted = m_Encryptor.Crypt(decryptedPacketBuffer.GetData());
            // Console.WriteLine("[WARDEN] Sending Encrypted:" + Environment.NewLine +
            //     DataFormatter.Format(m_bEncrypted));
            return m_bEncrypted;
        }

        /// <summary>
        /// Creates a new Warden module based on the seed provided.
        /// </summary>
        /// <param name="seed">The four-byte seed from the beginning of the CD-key hash.</param>
        public WardenModule(byte[] seed) {
            m_Random = new WardenRandom(seed);
            m_Encryptor = new SimpleCrypto(m_Random.GetBytes(0x10));
            m_Decryptor = new SimpleCrypto(m_Random.GetBytes(0x10));
        }
    }

    /// <summary>This is the XOR simple encryption class.</summary>
    internal class SimpleCrypto {
        private byte[] m_Key;
        /// <summary>Gets the encryption key.</summary>
        public byte[] Key { get { return m_Key; } }

        /// <summary>Generates the key based on the baseData</summary>
        /// <param name="baseData"></param>
        public SimpleCrypto(byte[] baseData) {
            char val = (char) 0;
            int position = 0;
            byte temp;

            m_Key = new byte[0x102];

            for (int i = 0; i < 0x100; i++)
                m_Key[i] = (byte) i;

            m_Key[0x100] = 0;
            m_Key[0x101] = 0;

            for (int i = 1; i <= 0x40; i++) {
                for (int j = 4; j > 0; j--) {
                    val += (char) (this.m_Key[(i * 4) - j] + baseData[position++ % baseData.Length]);
                    temp = this.m_Key[(i * 4) - j];
                    m_Key[(i * 4) - j] = m_Key[val & 0x0FF];
                    m_Key[val & 0x0FF] = temp;
                }
            }
        }

        /// <summary>Encrypts and decrypts data</summary>
        /// <param name="data"></param>
        /// <returns></returns>
        public byte[] Crypt(byte[] data) {
            int i;
            byte temp;

            for (i = 0; i < data.Length; i++) {
                m_Key[0x100]++;
                m_Key[0x101] += m_Key[m_Key[0x100] & 0x0FF];
                temp = m_Key[m_Key[0x101] & 0x0FF];
                m_Key[m_Key[0x101] & 0x0FF] = m_Key[m_Key[0x100] & 0x0FF];
                m_Key[m_Key[0x100] & 0x0FF] = temp;

                data[i] = (byte) (data[i] ^ m_Key[(m_Key[m_Key[0x101] & 0x0FF] + m_Key[m_Key[0x100] & 0x0FF]) & 0x0FF]);
            }

            return data;
        }
    }

    /// <summary>This structure is used to create the decryption base for the simple encryption class.</summary>
    internal class WardenRandom {
        private int m_Position;
        private byte[] m_Data1;
        private byte[] m_Data2;
        private byte[] m_Data3;
        private SHA1CryptoServiceProvider m_Sha1 = new SHA1CryptoServiceProvider();

        /// <summary>
        /// Creates a new WardenRandom class based on the seed provided.
        /// </summary>
        /// <param name="seed">The four-byte seed.</param>
        public WardenRandom(byte[] seed) {
            m_Data1 = new byte[0x14];
            m_Data2 = new byte[0x14];
            m_Data3 = new byte[0x14];

            int length1 = (int) seed.Length >> 1;   //2
            int length2 = seed.Length - length1;    //2

            byte[] seed1 = new byte[length1];
            byte[] seed2 = new byte[length2];

            for (int i = 0; i < length1; i++)
                seed1[i] = seed[i];
            for (int i = 0; i < length2; i++)
                seed2[i] = seed[i + length1];

            m_Data2 = m_Sha1.ComputeHash(seed1);
            m_Data3 = m_Sha1.ComputeHash(seed2);

            update();

            m_Position = 0;
        }

        private void update() {
            byte[] shathis = new byte[m_Data1.Length + m_Data2.Length + m_Data3.Length];
            m_Data2.CopyTo(shathis, 0);
            m_Data1.CopyTo(shathis, m_Data2.Length);
            m_Data3.CopyTo(shathis, m_Data2.Length + m_Data1.Length);
            m_Data1 = m_Sha1.ComputeHash(shathis);
        }

        /// <summary>
        /// Gets an array of bytes Warden style.
        /// </summary>
        /// <param name="a_iCount">The count of bytes.</param>
        /// <returns>A byte array of the hashed data.</returns>
        public byte[] GetBytes(int count) {
            byte[] bytes = new byte[count];
            for (int i = 0; i < count; i++)
                bytes[i] = GetByte();
            return bytes;
        }

        private byte GetByte() {
            int pos = m_Position;
            byte val = m_Data1[pos];
            pos++;
            if (pos >= 0x14) {
                pos = 0;
                update();
            }
            m_Position = pos;
            return val;
        }
    }
}
